#include <immintrin.h>
#include <windows.h>
#include <iostream>
#include <random>
using namespace std;
#define N 4096

float A[N][N];

void createMatrix() {
    srand(20000214);
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            A[i][j] = rand() % 100;
            // cout << A[i][j] << " ";
        }
        // cout << endl;
    }
    // cout << endl;
}

void LU() {
    for (int k = 0; k < N; k++) {
        for (int j = k + 1; j < N; j++) {
            A[k][j] /= A[k][k];
        }
        A[k][k] = 1;
        for (int i = k + 1; i < N; i++) {
            for (int j = k + 1; j < N; j++) {
                A[i][j] -= A[i][k] * A[k][j];
            }
            A[i][k] = 0;
        }
    }
}

void SSE_LU() {
    __m128 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        float kk[4] = {A[k][k], A[k][k], A[k][k], A[k][k]};
        t1 = _mm_loadu_ps(kk);
        for (int j = N - 4; j >= k; j -= 4) {
            t2 = _mm_loadu_ps(A[k] + j);
            t3 = _mm_div_ps(t2, t1);
            _mm_storeu_ps(A[k] + j, t3);
        }
        for (int j = k; j % 4 != N % 4; j++) {
            A[k][j] /= kk[0];
        }
        for (int i = k + 1; i < N; i++) {
            float ik[4] = {A[i][k], A[i][k], A[i][k],
                           A[i][k]};
            t1 = _mm_loadu_ps(ik);
            for (int j = N - 4; j > k; j -= 4) {
                t2 = _mm_loadu_ps(A[i] + j);
                t3 = _mm_loadu_ps(A[k] + j);
                t4 = _mm_sub_ps(t2, _mm_mul_ps(t1, t3));
                _mm_storeu_ps(A[i] + j, t4);
            }
            for (int j = k + 1; j % 4 != N % 4; j++) {
                A[i][j] -= A[i][k] * A[k][j];
            }
            A[i][k] = 0;
        }
    }
}

void SSE_LU_align() {
    __m128 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        __attribute__((aligned(16))) float kk[4] = {A[k][k], A[k][k],
                                                    A[k][k], A[k][k]};
        t1 = _mm_load_ps(kk);
        int j = k;
        while ((long long)(&A[k][j]) % 16 != 0) {
            j++;
        }
        for (int l = k; l < j; l++) {
            A[k][l] /= kk[0];
        }
        for (; j < N; j += 4) {
            t2 = _mm_load_ps(A[k] + j);
            t3 = _mm_div_ps(t2, t1);
            _mm_store_ps(A[k] + j, t3);
        }
        if (j > N) {
            for (int l = j - 4; l < N; l++) {
                A[k][l] /= kk[0];
            }
        }
        for (int i = k + 1; i < N; i++) {
            __attribute__((aligned(16))) float ik[4] = {
                A[i][k], A[i][k], A[i][k], A[i][k]};
            t1 = _mm_load_ps(ik);
            int j = k;
            while ((long long)(&A[i][j]) % 16 != 0) {
                j++;
            }
            for (int l = k; l < j; l++) {
                A[i][l] -= ik[0] * A[k][l];
            }
            for (; j < N; j += 4) {
                t2 = _mm_load_ps(A[i] + j);
                t3 = _mm_load_ps(A[k] + j);
                t4 = _mm_sub_ps(t2, _mm_mul_ps(t1, t3));
                _mm_store_ps(A[i] + j, t4);
            }
            if (j > N) {
                for (int l = j - 4; l < N; l++) {
                    A[i][l] -= ik[0] * A[k][l];
                }
            }
            A[i][k] = 0;
        }
    }
}

void AVX_LU() {
    __m256 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        float kk[8] = {A[k][k], A[k][k], A[k][k], A[k][k],
                       A[k][k], A[k][k], A[k][k], A[k][k]};
        t1 = _mm256_loadu_ps(kk);
        for (int j = N - 8; j >= k; j -= 8) {
            t2 = _mm256_loadu_ps(A[k] + j);
            t3 = _mm256_div_ps(t2, t1);
            _mm256_storeu_ps(A[k] + j, t3);
        }
        for (int j = k; j % 8 != N % 8; j++) {
            A[k][j] /= kk[0];
        }
        for (int i = k + 1; i < N; i++) {
            float ik[8] = {A[i][k], A[i][k], A[i][k],
                           A[i][k], A[i][k], A[i][k],
                           A[i][k], A[i][k]};
            t1 = _mm256_loadu_ps(ik);
            for (int j = N - 8; j > k; j -= 8) {
                t2 = _mm256_loadu_ps(A[i] + j);
                t3 = _mm256_loadu_ps(A[k] + j);
                t4 = _mm256_sub_ps(t2, _mm256_mul_ps(t1, t3));
                _mm256_storeu_ps(A[i] + j, t4);
            }
            for (int j = k + 1; j % 8 != N % 8; j++) {
                A[i][j] -= A[i][k] * A[k][j];
            }
            A[i][k] = 0;
        }
    }
}

void AVX_LU_align() {
    __m256 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        __attribute__((aligned(32))) float kk[8] = {
            A[k][k], A[k][k], A[k][k], A[k][k],
            A[k][k], A[k][k], A[k][k], A[k][k]};
        t1 = _mm256_load_ps(kk);
        int j = k;
        while ((long long)(&A[k][j]) % 32 != 0) {
            j++;
        }
        for (int l = k; l < j; l++) {
            A[k][l] /= kk[0];
        }
        for (; j < N; j += 8) {
            t2 = _mm256_load_ps(A[k] + j);
            t3 = _mm256_div_ps(t2, t1);
            _mm256_store_ps(A[k] + j, t3);
        }
        if (j > N) {
            for (int l = j - 8; l < N; l++) {
                A[k][l] /= kk[0];
            }
        }
        for (int i = k + 1; i < N; i++) {
            __attribute__((aligned(32))) float ik[8] = {
                A[i][k], A[i][k], A[i][k], A[i][k],
                A[i][k], A[i][k], A[i][k], A[i][k]};
            t1 = _mm256_load_ps(ik);
            int j = k;
            while ((long long)(&A[i][j]) % 32 != 0) {
                j++;
            }
            for (int l = k; l < j; l++) {
                A[i][l] -= ik[0] * A[k][l];
            }
            for (; j < N; j += 8) {
                t2 = _mm256_load_ps(A[i] + j);
                t3 = _mm256_load_ps(A[k] + j);
                t4 = _mm256_sub_ps(t2, _mm256_mul_ps(t1, t3));
                _mm256_store_ps(A[i] + j, t4);
            }
            if (j > N) {
                for (int l = j - 8; l < N; l++) {
                    A[i][l] -= ik[0] * A[k][l];
                }
            }
            A[i][k] = 0;
        }
    }
}

void SSE_LU_1() {
    __m128 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        float kk[4] = {A[k][k], A[k][k], A[k][k], A[k][k]};
        t1 = _mm_loadu_ps(kk);
        for (int j = N - 4; j >= k; j -= 4) {
            t2 = _mm_loadu_ps(A[k] + j);
            t3 = _mm_div_ps(t2, t1);
            _mm_storeu_ps(A[k] + j, t3);
        }
        for (int j = k; j % 4 != N % 4; j++) {
            A[k][j] /= kk[0];
        }
        for (int i = k + 1; i < N; i++) {
            for (int j = k + 1; j < N; j++) {
                A[i][j] -= A[i][k] * A[k][j];
            }
            A[i][k] = 0;
        }
    }
}

void SSE_LU_2() {
    __m128 t1, t2, t3, t4;
    for (int k = 0; k < N; k++) {
        for (int j = k + 1; j < N; j++) {
            A[k][j] /= A[k][k];
        }
        A[k][k] = 1;
        for (int i = k + 1; i < N; i++) {
            float ik[4] = {A[i][k], A[i][k], A[i][k],
                           A[i][k]};
            t1 = _mm_loadu_ps(ik);
            for (int j = N - 4; j > k; j -= 4) {
                t2 = _mm_loadu_ps(A[i] + j);
                t3 = _mm_loadu_ps(A[k] + j);
                t4 = _mm_sub_ps(t2, _mm_mul_ps(t1, t3));
                _mm_storeu_ps(A[i] + j, t4);
            }
            for (int j = k + 1; j % 4 != N % 4; j++) {
                A[i][j] -= A[i][k] * A[k][j];
            }
            A[i][k] = 0;
        }
    }
}

int main() {
    createMatrix();
    long freq, head, tail;
    QueryPerformanceFrequency((LARGE_INTEGER*)&freq);
    QueryPerformanceCounter((LARGE_INTEGER*)&head);
    LU();
    QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    cout << "serial: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // SSE_LU();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "SSE: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // SSE_LU_align();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "SSE align: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // AVX_LU();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "AVX: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // AVX_LU_align();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "AVX align: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    createMatrix();
    QueryPerformanceCounter((LARGE_INTEGER*)&head);
    SSE_LU_1();
    QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    cout << "SSE first: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    createMatrix();
    QueryPerformanceCounter((LARGE_INTEGER*)&head);
    SSE_LU_2();

    QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    cout << "SSE second: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    createMatrix();
    QueryPerformanceCounter((LARGE_INTEGER*)&head);
    SSE_LU();
    QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    cout << "SSE: " << (tail - head) * 1000.0 / freq << "ms" << endl;
}
